import logging
import os
import time
import tensorflow as tf
from tensorflow.contrib import slim
from tensorflow.contrib.tensorboard.plugins import projector
from tensorflow.python.framework import graph_util

import numpy as np

import cv2


def pad(total_boxes, w, h):

    tmpw = (total_boxes[:, 2] - total_boxes[:, 0] + 1).astype(np.int32)
    tmph = (total_boxes[:, 3] - total_boxes[:, 1] + 1).astype(np.int32)
    numbox = total_boxes.shape[0]

    dx = np.ones((numbox), dtype=np.int32)
    dy = np.ones((numbox), dtype=np.int32)
    edx = tmpw.copy().astype(np.int32)
    edy = tmph.copy().astype(np.int32)

    x = total_boxes[:, 0].copy().astype(np.int32)
    y = total_boxes[:, 1].copy().astype(np.int32)
    ex = total_boxes[:, 0].copy().astype(np.int32)
    ey = total_boxes[:, 1].copy().astype(np.int32)

    tmp = np.where(ex > w)
    edx.flat[tmp] = np.expand_dims(-ex[tmp] + w + tmpw[tmp], 1)
    ex[tmp] = w

    tmp = np.where(ey > h)
    edy.flat[tmp] = np.expand_dims(-ey[tmp] + h + tmph[tmp], 1)
    ey[tmp] = h

    tmp = np.where(x < 1)
    dx.flat[tmp] = np.expand_dims(2 - x[tmp], 1)
    x[tmp] = 1

    tmp = np.where(y < 1)
    dy.flat[tmp] = np.expand_dims(2 - y[tmp], 1)
    y[tmp] = 1

    return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph


# 填充为正方形
def rerec(bboxA):

    h = bboxA[:, 3] - bboxA[:, 1]
    w = bboxA[:, 2] - bboxA[:, 0]
    size = np.maximum(w, h)
    bboxA[:, 0] = bboxA[:, 0] + w * 0.5 - size * 0.5
    bboxA[:, 1] = bboxA[:, 1] + h * 0.5 - size * 0.5
    bboxA[:, 2:4] = bboxA[:, 0:2] + np.transpose(np.tile(size, (2, 1)))
    return bboxA

def bbreg(boundingbox, reg):


    if reg.shape[1] == 1:
        reg = np.reshape(reg, (reg.shape[2], reg.shape[3]))

    bb = boundingbox.copy()    

    w = boundingbox[:, 2] - boundingbox[:, 0] + 1

    h = boundingbox[:, 3] - boundingbox[:, 1] + 1


    bb[:,0] = boundingbox[:, 0] + reg[:, 0] * w
    bb[:,1] = boundingbox[:, 1] + reg[:, 1] * h
    bb[:,2] = boundingbox[:, 0] + reg[:, 2] * w
    bb[:,3] = boundingbox[:, 1] + reg[:, 3] * h

    # boundingbox[:, 0:4] = np.transpose(np.vstack([b1, b2, b3, b4]))
    
    return bb



def generateBoundingBox(imap, reg, scale, t):

    stride = 2
    cellsize = 12

    imap = np.transpose(imap)
    dx1 = np.transpose(reg[:, :, 0])
    dy1 = np.transpose(reg[:, :, 1])
    dx2 = np.transpose(reg[:, :, 2])
    dy2 = np.transpose(reg[:, :, 3])
    y, x = np.where(imap >= t)
    # print("yshape:" + str(y.shape))
    if y.shape[0] == 1:
        dx1 = np.flipud(dx1)
        dy1 = np.flipud(dy1)
        dx2 = np.flipud(dx2)
        dy2 = np.flipud(dy2)
    score = imap[(y, x)]
    reg = np.transpose(np.vstack([dx1[(y, x)], dy1[(y, x)],
                                  dx2[(y, x)], dy2[(y, x)]]))
    if reg.size == 0:
        reg = np.empty((0, 3))
    # bb-gridId, 满足score阈值条件    
    bb = np.transpose(np.vstack([y, x]))
    # print(bb)
    # q1为左上角的x1，y1
    q1 = np.fix((stride * bb + 1) / scale)
    # print(q1)
    # q2位右下角的x2，y2
    q2 = np.fix((stride * bb + cellsize) / scale)
    # print(q2)
    boundingbox = np.hstack([q1, q2, np.expand_dims(score, 1), reg])
    return boundingbox, reg

# def generateBoundingBox(imap, reg, scale, t):

#     stride = 2
#     cellsize = 12

#     imap = np.transpose(imap)
#     dx1 = np.transpose(reg[:,0])
#     dy1 = np.transpose(reg[:,1])
#     dx2 = np.transpose(reg[:,2])
#     dy2 = np.transpose(reg[:,3])
#     i = np.where(imap >= t)
#     print(i[1])
#     # if y.shape[0] == 1:
#     #     dx1 = np.flipud(dx1)
#     #     dy1 = np.flipud(dy1)
#     #     dx2 = np.flipud(dx2)
#     #     dy2 = np.flipud(dy2)
#     score = imap[i]
#     print(dx1.shape)
#     reg = np.transpose(np.vstack([dx1[i[1]], dy1[i[1]],
#                                   dx2[i[1]], dy2[i[1]]]))
#     if reg.size == 0:
#         reg = np.empty((0, 3))
#     bb = np.transpose(np.vstack(i))
#     q1 = np.fix((stride * bb + 1) / scale)
#     q2 = np.fix((stride * bb + cellsize - 1 + 1) / scale)
#     boundingbox = np.hstack([q1, q2, np.expand_dims(score, 1), reg])
#     return boundingbox, reg



def detect():

    im = cv2.imread('./test.jpg')
    im = im.astype('uint8')
    w,h,ch = im.shape
    print("import image shape:" + str(im.shape))
    
    # image = (im - 127.5) / 128.0
    #print(image)

    with tf.Session() as sess:
    # with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
        with open('./output_P/model.pb', "rb") as f:
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="")

            print("------------add pyramid-----------------")
            pyramid = []
            ww,hh,ch = im.shape
            minLenght = min(ww,hh)
            imPyramid = im.copy()
            minsize = 20
            m = 12.0 / minsize
            scales = []
            factor_counter = 0
            factor = 0.709
            ml = minLenght
            while ml >= 12 :
                # print("------------add pyramid-----------------")
                scaleFactor = m * np.power(factor, factor_counter)
                scales += [scaleFactor]
                # print("scaleFactor:" + str(scaleFactor))
                ww = w * scaleFactor
                hh = h * scaleFactor
                imPyramid = cv2.resize(im,(int(hh),int(ww)))
                pyramid.append((imPyramid.copy() - 127.5) / 128.0)
                # cv2.imwrite("./pyramid_"+str(imPyramid.shape[0]) + ".jpg", imPyramid)

                # print("ww:" +str(ww))
                factor_counter += 1
                scaleFactor = m * np.power(factor, factor_counter)
                ml = minLenght * scaleFactor
                # print("ml:" + str(ml))
                
                

            # pyramid.clear()
            # pyramid.append((im - 127.5) / 128.0)    


            boxes_all = np.empty((0,9))
            #for py in pyramid:
            print("------------PNet-----------------")
            for j in range(len(scales)):
                
                scale = scales[j]
                py = pyramid[j]
                # print("-------------Pyramid with " + str(py.shape) + "  is started-------------")
                feed_dict = {'input/shuffle_batch:0':[py]}
                _cls, _bb,_landmark = sess.run(['cls_fc/BiasAdd:0', 'bb_fc/BiasAdd:0','landmark_fc/BiasAdd:0'], feed_dict)

                # print("_cls shape:" + str(_cls.shape))
                softMax_ = np.exp(_cls[0,:,:,1]) / (np.exp(_cls[0,:,:,0]) + np.exp(_cls[0,:,:,1]))
                # print("max softmax:" + str(np.max(softMax_)))
                # print(softMax_.shape)

                boxes, _ = generateBoundingBox(softMax_,#_cls[0, :, :, 1],
                                       _bb[0, :, :, :],
                                       scale,
                                       0.8)
                # print("bounding boxes select : " + str(boxes.shape))                       
                if(len(boxes) > 0):
                    # print(boxes)
                    pick = nms(boxes.copy(), 0.5, 'Union')
                    # print("nms picked boxes:" + str(pick.shape))
                    # for p in pick:
                    #     print(boxes[p,:].shape)
                    boxes_all = np.append( boxes_all ,boxes[pick,:], axis = 0)

            print("boxes picked after PNet: " + str(boxes_all.shape))
            
            numbox = boxes_all.shape[0]
            if numbox > 0:
                pick = nms(boxes_all.copy(), 0.7, 'Union')
                boxes_all = boxes_all[pick, :]
                regw = boxes_all[:, 2] - boxes_all[:, 0]
                regh = boxes_all[:, 3] - boxes_all[:, 1]
                # 通过bb的回归值对bb对精调
                qq1 = boxes_all[:, 0] + boxes_all[:, 5] * regw
                qq2 = boxes_all[:, 1] + boxes_all[:, 6] * regh
                qq3 = boxes_all[:, 0] + boxes_all[:, 7] * regw
                qq4 = boxes_all[:, 1] + boxes_all[:, 8] * regh
                boxes_all = np.transpose(np.vstack([qq1, qq2, qq3, qq4,
                                                    boxes_all[:, 4]]))
                boxes_all = rerec(boxes_all.copy())
                boxes_all[:, 0:4] = np.fix(boxes_all[:, 0:4]).astype(np.int32)
                dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(
                    boxes_all.copy(), w, h)

                # print(boxes_all)
                im_bbr = im.copy()
                r = 0
                for box in boxes_all:
                    r += 1
                    # print(box)
                    cv2.rectangle(im_bbr,
                                (int(box[0]), int(box[1])),
                                (int(box[2]), int(box[3])), 
                                (0,255,0),1)

                # for k in range(len(x)):
                #     cv2.rectangle(im_bbr,
                #                     (x[k], y[k]),
                #                     (x[k]+tmpw[k], x[k]+tmph[k]), 
                #                     (0,0,255),1)
                print("------------Saving P Image-----------------")
                cv2.imwrite( "./test_bbr_P.jpg",im_bbr)      
                # cv2.imwrite( "./bbr_crop/test_bbr_R" + str(r) +".jpg",im_bbr[int(box[0]):(int(box[1])), 
                #                                                             int(box[2]):(int(box[3])),
                #                                                             :])           

            # 组成boundingBox和评分值
            # print(bb_pyramid)
           
    
    
        
        # with tf.Graph().as_default():
    
    tf.reset_default_graph()
    
   
    with tf.Graph().as_default() as graph:   
        with open('./output_R/model.pb', "rb") as f:
            output_graph_def = tf.GraphDef()
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="") 
    
            with tf.Session() as sess:
                # print(output_graph_def)
                batch = 128
                R_image_array_in = np.zeros((batch,24,24,3))
                v = 0
                # print("boxes_all shape:" + str(boxes_all.shape))
                # for k in range(0, len(boxes_all)):
                for box in boxes_all:
                    if v >= batch:
                        break

                    # tmp = np.zeros((int(tmph[k]), int(tmpw[k]), 3))
                    #根据生成的bb从原图上抠图
                    # tmp[dy[k] - 1:edy[k], dx[k] - 1:edx[k],:] = \
                    #     im[y[k] - 1:ey[k], x[k] - 1:ex[k], :]
                    # if (tmp.shape[0] > 0 and tmp.shape[1] > 0 or
                    #         tmp.shape[0] == 0 and tmp.shape[1] == 0):
                    #     R_image_array_in[k, :, :, :] = cv2.resize(tmp, (24, 24),interpolation=cv2.INTER_AREA)
                    # else:
                    #     return np.empty()

                    
                    _x = int(box[0])                    
                    _y = int(box[1])
                    _dx = int(box[2])
                    _dy = int(box[3])
                    try:
                        R_image_array_in[v, :, :, :] = cv2.resize(im[_y:_dy,_x:_dx:], (24, 24),interpolation=cv2.INTER_AREA)
                        cv2.imwrite( "./bbr_crop/test_bbr_R" + str(v) +".jpg",R_image_array_in[v,:,:,:]) 
                        # cv2.imwrite( "./bbr_crop/test_bbr_R" + str(v) +".jpg",im[_y:_dy,_x:_dx,:]) 
                        
                        v += 1
                    except:
                        pass
                    finally:
                        pass
                       
                   

                R_image_array_in = (R_image_array_in - 127.5) / 128.0

                # for i in range(R_image_array_in.shape[3]):

                # print("img:" + str(R_image_array_in.shape))
                # m = R_image_array_in[1:17,:,:,:]
                print("------------RNet-----------------")
                feed_dict = {'input/shuffle_batch:0':R_image_array_in}
                __cls, __bb,__landmark = sess.run(['cls_fc/BiasAdd:0', 'bb_fc/BiasAdd:0','landmark_fc/BiasAdd:0'], feed_dict)
                # print(__cls)
                softMax_ = np.exp(__cls[:,1]) / (np.exp(__cls[:,0]) + np.exp(__cls[:,1]))
                # print(softMax_)
                # print("bb shape:" + str(__bb.shape))
                # print("softMax_max:" + str(np.max(softMax_)))
                # print("softMax_ shape:" + str(softMax_.shape))

                ipass = np.where(softMax_ > 0.8)
                # print(ipass)
                #根据输出的softmax选出评分最高的bb
                total_boxes = boxes_all[ipass,:][0]

                # ipass_box为输出的bbReg里对应值
                if (len(ipass[0]) >0):
                    iPass_box = np.zeros([len(ipass[0]),5])
                    for j in range(len(ipass[0])):
                        ip = ipass[0][j]
                        # print(iPass_box[j][0:4])
                        # print(boxes_all[ip][0:4])
                        iPass_box[j][0:4] = __bb[ip][0:4]
                        iPass_box[j][4] = softMax_[ip]
                    
                    # print(iPass_box)

                    # print("total_boxes.shape" + str(total_boxes.shape))
                    pick = nms(total_boxes.copy(), 0.7, 'Union')
                    # print(pick)
                    #输入的源bb坐标

                    #用bb的回归量，将输入的bb再次调整
                    # print(total_boxes)
                    total_boxes = bbreg(total_boxes[pick,:].copy(), iPass_box[pick,:])
                    # print("输入的bb bbr")
                    # print(total_boxes)
                    # total_boxes = rerec(total_boxes.copy())
                    # print("输入的bb rerec")
                    # print(total_boxes)

                    im_bbr = im.copy()
                    
                    r = 0
                    for box in total_boxes:
                        # print("box")
                        # print(box)
                        im_bbr = cv2.rectangle(im_bbr,
                                    (int(box[0]),int(box[1])),
                                    (int(box[2]),int(box[3])), 
                                    (0,255,0),1)

                            

                        r+=1  
                    print("------------Saving RNet Image-----------------")
                    cv2.imwrite( "./test_bbr_R" +".jpg",im_bbr)        
                else:
                    print("no pick from Rnet-softmax")


    tf.reset_default_graph()
    with tf.Graph().as_default() as graph:   
        with open('./output_O/model.pb', "rb") as f:
            output_graph_def = tf.GraphDef()
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="") 
    
            with tf.Session() as sess:
                # print(output_graph_def)
                batch = 64
                O_image_array_in = np.zeros((batch,48,48,3))
                
                # print("boxes_all shape:" + str(total_boxes.shape))
                
                v = 0
                for box in total_boxes:
                    if(v >= batch):
                        break
                    _x =  int(box[0])                    
                    _y =  int(box[1])
                    _dx = int(box[2])
                    _dy = int(box[3])

                    try:
                        O_image_array_in[v, :, :, :] = cv2.resize(im[_y:_dy,_x:_dx:], (48, 48),interpolation=cv2.INTER_AREA)
                    except:
                        pass
                    finally:
                        pass    
                    v += 1  

                O_image_array_in = (O_image_array_in - 127.5) / 128.0


                feed_dict = {'input/shuffle_batch:0':O_image_array_in}
                print("------------ONet-----------------")
                __cls, __bb,__landmark = sess.run(['cls_fc/BiasAdd:0', 'bb_fc/BiasAdd:0','landmark_fc/BiasAdd:0'], feed_dict)  

                # print(__landmark.shape) 
                # print(__landmark) 
                softMax_ = np.exp(__cls[:,1]) / (np.exp(__cls[:,0]) + np.exp(__cls[:,1]))
                # print(softMax_)
                ipass = np.where(softMax_ > 0.95)
                # print(ipass)
                total_boxes = total_boxes[ipass,:][0]
                # print(total_boxes)
                __bb = __bb[ipass,:][0]
                # print(__bb)
                __landmark = __landmark[ipass,:][0]
                # print(__landmark)

                pick = nms(total_boxes.copy(), 0.7, 'Min')
                # print(pick)
                #输入的源bb坐标

                #用bb的回归量，将输入的bb再次调整
                # print(total_boxes)
                total_boxes = bbreg(total_boxes[pick,:].copy(), iPass_box[pick,:])

                v = 0
                im_bbr_o = im.copy()
                for box in total_boxes:
                    im_bbr_o = cv2.rectangle(im_bbr_o,
                                    (int(box[0]),int(box[1])),
                                    (int(box[2]),int(box[3])), 
                                    (0,255,0),1)
                    land = __landmark[v]
                    land[[0,2,4,6,8]] = box[0] + land[[0,2,4,6,8]] * (box[2] - box[0])
                    land[[1,3,5,7,9]] = box[1] + land[[1,3,5,7,9]] * (box[3] - box[1])
                    # print(land)
                    for i in range(5):
                        im_bbr_o = cv2.rectangle(im_bbr_o,
                                        (int(land[0 + i * 2]),int(land[1 + i * 2])),
                                        (int(land[0 + i * 2] + 2),int(land[1 + i * 2] + 2)), 
                                        (0,255,0),1)


                    v += 1
                cv2.imwrite( "./test_bbr_O"+".jpg",im_bbr_o)   


def nms(boxes, threshold, method):

    if boxes.size == 0:
        return np.empty((0, 3))
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    s = boxes[:, 4]
    # x1 = boxes[:, 0]
    # y1 = boxes[:, 1]
    # x2 = boxes[:, 2] + boxes[:, 0]
    # y2 = boxes[:, 3] + boxes[:, 1]
    # s = score[:, 1]
    
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    s_sort = np.argsort(s)
    pick = np.zeros_like(s, dtype=np.int16)
    counter = 0
    while s_sort.size > 0:
        i = s_sort[-1]
        pick[counter] = i
        counter += 1
        idx = s_sort[0:-1]
        xx1 = np.maximum(x1[i], x1[idx])
        yy1 = np.maximum(y1[i], y1[idx])
        xx2 = np.minimum(x2[i], x2[idx])
        yy2 = np.minimum(y2[i], y2[idx])
        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        if method is 'Min':
            o = inter / np.minimum(area[i], area[idx])
        else:
            o = inter / (area[i] + area[idx] - inter)
        s_sort = s_sort[np.where(o <= threshold)]
    pick = pick[0:counter]
    return pick

if __name__ == '__main__':
    detect()
